import streamlit as st
import cv2
import os
import time
from ultralytics import YOLO
from utils.model_utils import find_latest_best_model, clean_old_files
import zipfile

def create_zip_file(image_paths, zip_filename):
    """
    创建一个包含所有图像的 ZIP 文件。

    :param image_paths: 图像路径列表
    :param zip_filename: ZIP 文件名
    """
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for image_path in image_paths:
            zipf.write(image_path, os.path.basename(image_path))

def main():
    st.title("多SAR图识别")

    # 上传图像
    uploaded_files = st.file_uploader("选择SAR图像...", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
    if uploaded_files:
        try:
            # 生成带有时间戳的上传目录
            timestamp = int(time.time())
            upload_dir = os.path.join("uploads", str(timestamp))
            os.makedirs(upload_dir, exist_ok=True)
            image_paths = []

            for file in uploaded_files:
                image_filename = file.name
                image_path = os.path.join(upload_dir, image_filename)
                with open(image_path, "wb") as f:
                    f.write(file.read())
                image_paths.append(image_path)

            # 选择模型
            model_options = {
                "YOLOv11n": "yolo11n_best.pt",
                "YOLOv11s": "yolo11s_best.pt",
                "YOLOv11m": "yolo11m_best.pt",
                "YOLOv11l": "yolo11l_best.pt",
                "YOLOv11x": "yolo11x_best.pt"
            }
            selected_model_name = st.selectbox("选择YOLOv11模型", list(model_options.keys()))

            # 设置置信度阈值
            confidence_threshold = st.slider("置信度阈值", 0.0, 1.0, 0.5, 0.05)

            if st.button("开始检测"):
                # 查找最新的 best 模型文件
                current_dir = os.getcwd()
                parent_dir = os.path.dirname(current_dir)
                models_dir = os.path.join(parent_dir, "models", "v2")
                selected_model_path = os.path.join(models_dir, model_options[selected_model_name])

                if not os.path.exists(selected_model_path):
                    # 如果找不到选定的模型文件，使用最新的最佳模型
                    selected_model_path = find_latest_best_model(models_dir)
                    st.warning(f"找不到选定的模型文件，使用最新的最佳模型: {selected_model_path}")

                st.write(f"选择的模型文件: {selected_model_path}")
                # 创建保存检测结果的目录
                save_dir = os.path.join(current_dir, "runs", "detect", "predict")
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)

                # 加载模型
                model = YOLO(selected_model_path)

                # 处理每一张图像
                output_dir = "output"
                os.makedirs(output_dir, exist_ok=True)
                processed_images = []

                for image_path in image_paths:
                    # 使用模型预测
                    with st.spinner("正在检测..."):
                        results = model.predict(source=image_path, save=True, imgsz=1024, conf=confidence_threshold, save_dir=save_dir)
                    
                    # 从结果中获取保存的目录
                    save_path = results[0].save_dir  # 假设results[0]是第一个预测结果
                    if save_path:
                        for filename in os.listdir(save_path):
                            if filename.endswith(('.png', '.jpg', '.jpeg')):
                                image_path = os.path.join(save_path, filename)
                                
                                # 检查文件名是否已经存在于 processed_images 列表中
                                if not any(filename == existing_filename for existing_filename, _ in processed_images):
                                    processed_images.append((filename, image_path))

                # 显示所有处理后的图像
                images_to_show = processed_images[:4]

                # 创建4列布局
                cols = st.columns(min(4, len(images_to_show)))

                # 在4列中分别显示图像，并添加点击放大功能
                for i, (filename, path) in enumerate(images_to_show):
                    with cols[i]:
                        st.image(path, caption=f"检测到的飞机 - {filename}", width=280)

                # 初始化会话状态
                if 'downloaded_files' not in st.session_state:
                    st.session_state.downloaded_files = set()

                # 批量下载按钮
                zip_filename = "all_predictions.zip"
                create_zip_file([path for _, path in processed_images], zip_filename)
                with open(zip_filename, "rb") as zipf:
                    # 为 download_button 提供一个唯一的 key
                    st.download_button(
                        label="批量下载所有预测结果",
                        data=zipf,
                        file_name=zip_filename,
                        mime="application/zip",
                        key="download_all_predictions"  # 添加唯一的 key
                    )

                # 清理20次以前的临时文件
                clean_old_files(upload_dir, max_keep=20)

        except Exception as e:
            st.error(f"发生错误: {e}")

if __name__ == "__main__":
    main()